from typing import Callable

from pynumpak.type import func_t
from .solve import *


def least_square(a: Matrix, b: Vector, sigma: Callable = None):
    """a 的每一行是待拟合的自变量, b 是列向量, 代表因变量. sigma 用来处理异方差性"""

    if sigma is not None:
        n = len(b)
        c = Vector(n)
        for i in range(n):
            c[i] = sigma(a[:, i])
            a[:, i] /= c[i]
            b[i] /= c[i]
    return as_vec(solve(a * a.transpose, a * b.v_vec))


def poly_fit(a: Vector, b: Vector, n: int = 3, sigma: func_t = None):
    """a 是行向量, 表示自变量, b 是行向量, 代表因变量, 用 n 次多项式拟合, 输出多项式系数"""

    A = Matrix(n + 1, len(a))
    A[0] = 1
    for i in range(1, n + 1):
        A[i] = (a ** i).store_

    return least_square(A, b, sigma=(lambda x: sigma(x[1])) if sigma is not None else None)
